import cv2
import os

# 加载训练集
recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.read("./data/train.yml")

# 将id与姓名进行绑定(后面可以检索出该用户的名称)
def name():
    names = {}
    filePaths = [os.path.join('./train/',f) for f in os.listdir('./train/')]
    for fpath in filePaths:
        id = str(os.path.split(fpath)[1].split(".")[0])
        name = str(os.path.split(fpath)[1].split(".")[1])
        names[id] = name
    return names

#检测人脸
def faceDetection(frame):
    classifier = cv2.CascadeClassifier("./haarcascades/haarcascade_frontalface_alt2.xml")
    gray_img = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
    faces = classifier.detectMultiScale(gray_img,
                                1.01,
                                5,
                                0,
                                (100,100),
                                (500,500))
    names = name()
    for face in faces:
        (x,y,w,h) = face
        cv2.rectangle(frame,(x,y),(x+w,y+h),color=(255,0,0),thickness=2)
        # 使用训练好的数据集预测人脸数据
        id,confidence = recognizer.predict(gray_img[y:y+h,x:x+w])
        # confidence 置信评分，表明与实际结果的差距，差距越大，证明越不相似。
        if confidence < 80:
            cv2.putText(frame,names[str(id)],(x,y),0,1.2,color=(255,0,0),thickness=2)
        else:
            cv2.putText(frame, "unknown", (x, y),0,1.2, color=(255, 0, 0), thickness=2)

    cv2.imshow("frame",frame)

if __name__ == '__main__':
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    while True:
        fig,frame = cap.read()
        if fig is None:
            print("Capture No Found")
            break
        if cv2.waitKey(1) == ord(" "):
            break
        faceDetection(frame)
    cap.release()
    cv2.destroyAllWindows()